package com.springframework.boot.awesome.annotation.agent;

import com.springframework.boot.awesome.annotation.arch.exception.AgentLocationInvalidException;
import com.springframework.boot.awesome.annotation.arch.exception.AgentPropertiesException;
import javassist.*;
import org.apache.commons.lang3.StringUtils;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.*;

public class AgentClassFileTransformer implements ClassFileTransformer {

    public String targetClassPackage;
    public String targetClassMethod;
    public String targetClassMethodParams;
    public String agentCode;
    public AgentType agentLocation;

    public AgentClassFileTransformer(String targetClassPackage, String targetClassMethod,
                                     String targetClassMethodParams, String agentCode, AgentType agentLocation) {
        this.targetClassPackage = targetClassPackage;
        this.targetClassMethod = targetClassMethod;
        this.targetClassMethodParams = targetClassMethodParams;
        this.agentCode = agentCode;
        this.agentLocation = agentLocation;
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        if(StringUtils.isEmpty(targetClassPackage)){
            throw new IllegalClassFormatException("AgentProperties.targetClassPackage is null or \"\"");
        }
        if(StringUtils.isEmpty(targetClassMethod)) {
            throw new IllegalClassFormatException("AgentProperties.targetClassMethod is null or \"\"");
        }
        if(StringUtils.isEmpty(targetClassMethodParams)){
            throw new IllegalClassFormatException("AgentProperties.targetClassMethodParams is null or \"\"");
        }
        if(StringUtils.isEmpty(agentCode)){
            throw new IllegalClassFormatException("AgentProperties.agentCode is null or \"\"");
        }
        if(agentLocation == null){
            throw new IllegalClassFormatException("AgentProperties.agentCode is null or \"\"");
        }

        if(! targetClassPackage.equals(className)){
            return null;
        }

        //javassist的包名是用点分隔的，需要转换
        if(className != null && className.indexOf("/") != -1){
            className = className.replaceAll("/", ".");
        }

        try {
            //通过包名获取类文件
            CtClass ctClass = ClassPool.getDefault().get(className);
            //构造CtClass方法参数数组
            String[] paramsArgs = targetClassMethodParams.split(",");
            CtClass[] CtClassArgs = new CtClass[paramsArgs.length];
            for(int i = 0; i < paramsArgs.length; i ++){
                CtClassArgs[i] = ClassPool.getDefault().get(paramsArgs[i]);
            }

            //获得指定方法名的方法
            CtMethod doSomeThing = ctClass.getDeclaredMethod(targetClassMethod, CtClassArgs);

            if(agentLocation == AgentType.BEFORE){
                doSomeThing.insertBefore(agentCode);
            }else if(agentLocation == AgentType.AFTER){
                doSomeThing.insertAfter(agentCode);
            }else{
                System.out.println("agent.location properties is invalid");
                throw new AgentLocationInvalidException("agent.location properties is invalid, should: (before) or (after)");
            }

            return ctClass.toBytecode();
        } catch (NotFoundException e) {
            e.printStackTrace();
        } catch (CannotCompileException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (AgentLocationInvalidException e) {
            e.printStackTrace();
        }
        return null;
    }
}
